package org.neu.ccs.secureim;

import java.io.Serializable;
import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.Random;

import javax.crypto.KeyAgreement;
import javax.crypto.spec.DHParameterSpec;

public class DHValues implements Serializable {
	/**
	 * 
	 */
	private static final long serialVersionUID = 1L;
	PublicKey pubKey;
	PrivateKey privKey;
	DHParameterSpec spec;
	
	public DHValues(Random r) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException
	{
	    KeyPairGenerator kg = KeyPairGenerator.getInstance("DH");
	    BigInteger safePrime = GenerateSafePrime(r, 1024);
	    spec = new DHParameterSpec(safePrime, GetGenerator(safePrime), 1024);
	    kg.initialize(spec);
	    KeyPair pair = kg.generateKeyPair();

	    privKey = pair.getPrivate();
	    pubKey = pair.getPublic();
	}
	
	public DHValues(DHParameterSpec spec) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException
	{
		KeyPairGenerator kg = KeyPairGenerator.getInstance("DH");
		kg.initialize(spec);
	    KeyPair pair = kg.generateKeyPair();

	    privKey = pair.getPrivate();
	    pubKey = pair.getPublic();
	}
	
	public Key doFinal(PublicKey pubKey, String symAlg) throws InvalidKeyException, IllegalStateException, NoSuchAlgorithmException
	{
		KeyAgreement ka = KeyAgreement.getInstance("DH");
	    ka.init(privKey);
	    ka.doPhase(pubKey, true);

	    return ka.generateSecret(symAlg);
	}
	
	public PublicKey getPublicKey()
	{
		return pubKey;
	}
	
	public DHParameterSpec getParameterSpec()
	{
		return spec;
	}
	
	public static BigInteger GenerateSafePrime(Random r, int size)
	{		
		return new BigInteger(size, 10, r);
		/*while( true )
		{
			BigInteger a = new BigInteger(2047,10,r);
			BigInteger b = a.add(a).add(BigInteger.ONE);
			if( b.isProbablePrime(10) )
			{
				System.out.println("FOUND!");
				return b;
			}
		}*/
	}
	
	public static BigInteger GetGenerator(BigInteger safePrime)
	{
		return new BigInteger("5");
	}
}
